#include "PSO.h"

PSOOptimizer::PSOOptimizer(PSOPara* pso_para, ComputeFitness fitness_fun)
{
	particle_num_ = pso_para->particle_num_;
	max_iter_num_ = pso_para->max_iter_num_;
	dim_ = pso_para->dim_;
	curr_iter_ = 0;

	dt_ = new int[dim_];
	wstart_ = new double[dim_];
	wend_ = new double[dim_];
	C1_ = new double[dim_];
	C2_ = new double[dim_];

	for (int i = 0; i < dim_; i++) {
		dt_[i] = pso_para->dt_[i];
		wstart_[i] = pso_para->wstart_[i];
		wend_[i] = pso_para->wend_[i];
		C1_[i] = pso_para->C1_[i];
		C2_[i] = pso_para->C2_[i];
	}

	if (pso_para->upper_bound_ && pso_para->lower_bound_) {
		upper_bound_ = new int[dim_];
		lower_bound_ = new int[dim_];
		range_interval_ = new int[dim_];

		for (int i = 0; i < dim_; i++) {
			upper_bound_[i] = pso_para->upper_bound_[i];
			lower_bound_[i] = pso_para->lower_bound_[i];
			range_interval_[i] = upper_bound_[i] - lower_bound_[i];
		}
	}

	particles_ = new Particle[particle_num_];
    for(int i=0; i<particle_num_; i++) {
        particles_[i].dim_ = pso_para->dim_;
    }

	w_ = new double[dim_];
	all_best_position_ = new int[dim_];

	results_dim_ = pso_para->results_dim_;
	if (results_dim_) {
		results_ = new int[results_dim_];
	}
    rand_max = 550000;
	fitness_fun_ = fitness_fun;
}

PSOOptimizer::~PSOOptimizer()
{
	delete[]particles_;
	delete[]upper_bound_;
	delete[]lower_bound_;
	delete[]range_interval_;
	delete[]dt_;
	delete[]wstart_;
	delete[]wend_;
	delete[]w_;
	delete[]C1_;
	delete[]C2_;
	delete[]all_best_position_;
	delete[]results_;
}

void PSOOptimizer::InitialAllParticles(const vector<int>& init_position)
{
    // 初始化第一个粒子参数并设置最优值
	InitialParticle(0, init_position);
	all_best_fitness_ = particles_[0].best_fitness_;
	for (int j = 0; j < dim_; j++) {
		all_best_position_[j] = particles_[0].best_position_[j];
	}
    // 初始化其他粒子，并更新最优值
	for (int i = 1; i < particle_num_; i++) {
		InitialParticle(i, init_position);
#ifdef MAXIMIZE_FITNESS
		if (particles_[i].best_fitness_ > all_best_fitness_)
#else
		if (particles_[i].best_fitness_ < all_best_fitness_)
#endif
		{
			all_best_fitness_ = particles_[i].best_fitness_;
			for (int j = 0; j < dim_; j++) {
				all_best_position_[j] = particles_[i].best_position_[j];
			}

			if (particles_[i].results_dim_ && results_dim_ == particles_[i].results_dim_) {
				for (int k = 0; k < results_dim_; k++) {
					results_[k] = particles_[i].results_[k];
				}
			} else if (results_dim_) {
				std::cout << "WARNING: the dimension of your saved results for every particle\nis not match with the dimension you specified for PSO optimizer ant no result is saved!" << std::endl;
			}
		}
	}
}

void PSOOptimizer::InitialParticle(int i, vector<int> init_position)
{
    particles_[i].position_ = new int[dim_];
    particles_[i].velocity_ = new int[dim_];
    particles_[i].best_position_ = new int[dim_];
    for(int k=0; k<dim_; k++) {
//        particles_[i].position_[k] = 0;
//        particles_[i].velocity_[k] = 100;
//        particles_[i].best_position_[k] = init_position[k];
        particles_[i].position_[k] = init_position[k];
        particles_[i].velocity_[k] = abs(GetIntRand(100));
        particles_[i].best_position_[k] = abs(init_position[k]+(50-GetIntRand(100)));
    }

    particles_[i].fitness_ = GetFitness(particles_[i]);

    for (int j = 0; j < dim_; j++) {
        particles_[i].best_position_[j] = particles_[i].position_[j];
    }
    particles_[i].best_fitness_ = particles_[i].fitness_;
}

int PSOOptimizer::GetIntRand(int N)
{
    return rand() % (N + 1);
}


double PSOOptimizer::GetDoubleRand(int N)
{
	double temp = rand() % (N + 1) / (double)(N + 1);
	return temp;
}

long long PSOOptimizer::GetFitness(Particle & particle)
{
	return fitness_fun_(particle);
}

void PSOOptimizer::UpdateAllParticles()
{
	GetInertialWeight();
	for (int i = 0; i < particle_num_; i++) {
		UpdateParticle(i);
#ifdef MAXIMIZE_FITNESS
		if (particles_[i].best_fitness_ > all_best_fitness_)
#else
		if (particles_[i].best_fitness_ < all_best_fitness_)
#endif
		{
			all_best_fitness_ = particles_[i].best_fitness_;
			for (int j = 0; j < dim_; j++) {
				all_best_position_[j] = particles_[i].best_position_[j];
			}

            // 如果需要保存出一些结果
			if (particles_[i].results_dim_ && results_dim_ == particles_[i].results_dim_) {
				for (int k = 0; k < results_dim_; k++) {
					results_[k] = particles_[i].results_[k];
				}
			} else if (results_dim_) {
				std::cout << "WARNING: the dimension of your saved results for every particle\nis not match with the dimension you specified for PSO optimizer ant no result is saved!" << std::endl;
			}
		}
	}
	curr_iter_++;
}

void PSOOptimizer::UpdateParticle(int i)
{
	for (int j = 0; j < dim_; j++) {
		//double last_velocity = particles_[i].velocity_[j];
		int last_position = particles_[i].position_[j];

		particles_[i].velocity_[j] = w_[j] * particles_[i].velocity_[j] +
			C1_[j] * GetDoubleRand() * (particles_[i].best_position_[j] - particles_[i].position_[j]) +
			C2_[j] * GetDoubleRand() * (all_best_position_[j] - particles_[i].position_[j]);
		particles_[i].position_[j] += dt_[j] * particles_[i].velocity_[j];

		if (upper_bound_ && lower_bound_) {
			if (particles_[i].position_[j] > upper_bound_[j]) {
				double thre = GetDoubleRand(99);
				if (last_position == upper_bound_[j]) {
					particles_[i].position_[j] = GetDoubleRand() * range_interval_[j] + lower_bound_[j];
				} else if (thre < 0.5) {
					particles_[i].position_[j] = upper_bound_[j] - (upper_bound_[j] - last_position) * GetDoubleRand();
				} else {
					particles_[i].position_[j] = upper_bound_[j];
				}		
			}
			if (particles_[i].position_[j] < lower_bound_[j]) {
				double thre = GetDoubleRand(99);
				if (last_position == lower_bound_[j]) {
					particles_[i].position_[j] = GetDoubleRand() * range_interval_[j] + lower_bound_[j];
				} else if (thre < 0.5) {
					particles_[i].position_[j] = lower_bound_[j] + (last_position - lower_bound_[j]) * GetDoubleRand();
				} else {
					particles_[i].position_[j] = lower_bound_[j];
				}
			}
		}

	}
	particles_[i].fitness_ = GetFitness(particles_[i]);

#ifdef MAXIMIZE_FITNESS
	if (particles_[i].fitness_ > particles_[i].best_fitness_)
#else
	if (particles_[i].fitness_ < particles_[i].best_fitness_)
#endif
	{
		particles_[i].best_fitness_ = particles_[i].fitness_;
		for (int j = 0; j < dim_; j++) {
			particles_[i].best_position_[j] = particles_[i].position_[j];
		}
	}
}

void PSOOptimizer::GetInertialWeight()
{
	double temp = curr_iter_ / (double)max_iter_num_;
	temp *= temp;
	for (int i = 0; i < dim_; i++) {
		w_[i] = wstart_[i] - (wstart_[i] - wend_[i]) * temp;
	}
}

Particle::Particle(int dim, int * position, int * velocity, int * best_position, long long best_fitness)
{
	dim_ = dim;
	position_ = position;
	velocity_ = velocity;
	best_position_ = best_position;
	best_fitness_ = best_fitness;
}


long long FitnessFunction(Particle& particle)
{
    // position 传给 table1d
    int x_dim = particle.dim_;
    for(int i=0; i<x_dim; i++) {
        pso_solver.table1d[i] = int(particle.position_[i]);
    }
    // table1d转table
    pso_solver.table1d2table();
    // 计算fitness
    return pso_solver.get_cost();
}
